from kuai.ffio import *
from kuai.ff import *

PPF_HEADER = """#DFF:PPF
#PROTOCOL = %s
"""

class DffPPFReader(ForceFieldReader):
    def read(self, file):
        for i in file:
            if i.startswith("#DFF:PPF"):
                break
        i = file.next()
        assert i.startswith("#PROTOCOL")
        type = i.split("=")[1].strip()
        result = ForceField(type)
        for i in file:
            tokens = [j.strip() for j in i.split(":")]
            if len(tokens) >= 2:
                term = tokens[0].strip()
                atoms = [j.strip() for j in tokens[1].split(",")]
                key = ForceFieldKey(term, atoms)
                par = ForceFieldItem.parse(i)
                result[key] = par
            else:
                break
        return result 

reader = DffPPFReader()

class DffPPFWriter(ForceFieldWriter):
    def write(self, file, v):
        file.write(PPF_HEADER % v.type)
        keyset = [i for i in v.keys()]
        keyset.sort()
        for key in keyset:
            file.write(str(key)+str(v[key])+"\n")
            
writer = DffPPFWriter()
